package com.migo.utils;


import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * redis ip 限流
 *
 * @author wes.wang
 * @date 2022/11/24 22:28
 **/
@Component
public class RedisLimit {

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 是否允许操作
     * @param userId 用户ID（可视情况替换为ip等）
     * @param actionKey 操作行为的key
     * @param period 期限（单位：毫秒），如10秒内，则为10000
     * @param maxCount 最大数
     * @return
     */
    public boolean isActionAllowed(String userId, String actionKey, long period, int maxCount){
        // 这里的key是用来组合行为规则的，谁做了什么操作，比如用户登录操作限流，IP请求限流等，根据业务需求定义即可
        String key = actionKey + "_" + userId;
        byte[] keyByte = key.getBytes();

        // 获取到当前的时间戳，以便后面存入set集合中，判断是否在指定时间范围内
        long timestamp = System.currentTimeMillis();
        byte[] timestampByte = String.valueOf(timestamp).getBytes();

        // 进行一组redis操作，记录本次请求，并且删除已过期的请求，最后返回指定时间内的请求数
        List<Object> as = redisTemplate.executePipelined((RedisCallback<?>) c -> {
            // 记录行为时间
            c.zAdd(keyByte, timestamp, timestampByte);
            // 删除期限外的记录，则剩余的都是期限内的记录
            c.zRemRangeByScore(keyByte,0L, timestamp - period);
            // 获取期限内的行为数量
            c.zCard(keyByte);
            // 设置过期时间，避免冷用户持续占用内存，宽限1秒
            c.expire(keyByte, period/1000 + 1L);
            return null;
        });

        // 最后，判断指定时间段内的请求数是否达到指定的最大请求数，如果是则返回false，否则返回true
        long size = 0;
        Object sizeObj= as.get(2);
        if(sizeObj != null){
            size = (long) sizeObj;
        }
        return size <= maxCount;
    }
}
